import os
from collections import defaultdict
from PIL import Image
import matplotlib.pyplot as plt
import shutil
from tqdm import tqdm

def read_image(img_path):
    """Keep reading image until succeed.
    This can avoid IOError incurred by heavy IO process."""
    got_img = False
    if not os.path.exists(img_path):
        raise IOError("{} does not exist".format(img_path))
    while not got_img:
        try:
            img = Image.open(img_path).convert('RGB')
            got_img = True
        except IOError:
            print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
            pass
    return img

data_dir = '../../data/train'
pseudo_label_file = '/pse_label.txt'
save_path = data_dir + '/pse_label_image'
pseudo_label = open(data_dir + pseudo_label_file, 'r')
count_image = defaultdict(list)
while True:
    lines = pseudo_label.readline()
    if not lines:
        break
    (img_name, img_label) = lines.strip().split(':')
    count_image[img_label].append(img_name)
for ID in tqdm(count_image.keys()):
    if not os.path.exists(save_path+'/'+str(ID)):
        os.makedirs(save_path+'/'+str(ID))
    for img_name in count_image[ID]:
        shutil.copy(data_dir+'/images/'+img_name,save_path+'/'+str(ID))
        # image = read_image(data_dir+'/images/'+img_name)
